//
//  segmentTree.swift
//  SegmentTree
//
//  Created by lhc on 2024/8/16.
//

import Foundation

class SegmentTree<E> {
    // 线段树内部的树也是由一个我们填充后的满二叉树来体现的
    // 实现方式是一个数组(数组可以来实现完全二叉树和满二叉树)
    
    private var data: [E]
    private var tree: [E?]
    private var size: Int
    private var merge: (E?, E?) -> (E)
    
    init(array: [E], merge: @escaping (E?, E?) -> (E)) {
        self.data = array
        self.size = array.count
        
        // 生成的树默认为data长度的四倍
        // 因为二叉树最后一行大小约为之前的两倍
        // 线段树叶子结点不一定在最后一行，所以需要额外一行额外存储,没有值的内容由nil填充
        self.tree = Array<E?>(repeating: nil, count: self.size * 4)
        self.merge = merge
        
        self.buildSegmentTree()
    }
    
    /// 创建树结构
    private func buildSegmentTree() {
        self.buildSegmentTreeByIndex(0, left: 0, right: self.size - 1)
    }
    
    private func buildSegmentTreeByIndex(_ index: Int, left: Int, right: Int) {
        if left == right {
            // 递归到底 left == right了
            self.tree[index] = self.data[left]
            return
        }
        
        // 如果left = right那么说明现在存的不是叶子结点是一个区间，那么就有左孩子和右孩子
        let leftChildIndex = self.leftChildIndexBy(index: index)
        let rightChildIndex = self.rightChildIndexBy(index: index)
        let mid = (right - left) / 2 + left
        
        // 递归生成左孩子和右孩子
        self.buildSegmentTreeByIndex(leftChildIndex, left: left, right: mid)
        self.buildSegmentTreeByIndex(rightChildIndex, left: mid + 1, right: right)
        
        // 把左孩子和右孩子的值做一个merge赋值给当前结点
        self.tree[index] = self.merge(self.tree[leftChildIndex], self.tree[rightChildIndex])
    }
    
    /// 获取对应index的元素
    ///
    ///  获取的是原始数组的index的元素
    func getByIndex(_ index: Int) -> E? {
        if index < 0 || index >= self.size {
            return nil
        }
        
        return self.data[index]
    }
    
    /// 查询data[queryleft...queryright]内的元素的merge结果
    func query(queryLeft: Int, queryRight: Int) -> E? {
        guard queryLeft >= 0, queryLeft < self.size, queryRight >= 0, queryRight < self.size, queryLeft <= queryRight else {
            return nil
        }
        
        return self.queryByTreeIndex(0, left: 0, right: self.size - 1, queryLeft: queryLeft, queryRight: queryRight)
    }
    
    /// 递归查询data[queryleft...queryright]内的元素的merge结果
    ///
    ///  - treeindex：self.tree数组对应index的
    ///  - left: treeIndex对应的self.data[left...right]中的left
    ///  - right: treeIndex对应的self.data[left...right]中的right
    ///  - queryLeft: 查询的data[queryleft...queryright]中的queryleft
    ///  - queryRight 查询的data[queryleft...queryright]中的queryright
    private func queryByTreeIndex(_ treeIndex: Int, left: Int, right: Int, queryLeft: Int, queryRight: Int) -> E? {
        if left == queryLeft && right == queryRight {
            // 如果当前左边界右边界就是要查询的左右边界,那返回当前treeindex保存的值
            return self.tree[treeIndex]
        }
        
        
        let leftChildIndex = self.leftChildIndexBy(index: treeIndex)
        let rightChildIndex = self.rightChildIndexBy(index: treeIndex)
        
        let mid = (right - left) / 2 + left
        
        if queryLeft > mid {
            // 如果查询的左边界大于当前的中间值说明查询的内容都在当前的右子树中
            return self.queryByTreeIndex(rightChildIndex, left: mid + 1, right: right, queryLeft: queryLeft, queryRight: queryRight)
        }else if queryRight <= mid {
            // 如果查询的右边界小于当前的中间值说明查询的内容都在当前的左子树中
            return self.queryByTreeIndex(leftChildIndex, left: left, right: mid, queryLeft: queryLeft, queryRight: queryRight)
        }else {
            // 到这里说明要查询的边界有在左子树的也有在右子树的
            
            let leftResult = self.queryByTreeIndex(leftChildIndex, left: left, right: mid, queryLeft: queryLeft, queryRight: mid)
            let rightResult = self.queryByTreeIndex(rightChildIndex, left: mid + 1, right: right, queryLeft: mid + 1, queryRight: queryRight)
            
            return self.merge(leftResult, rightResult)
        }
        
        
    }
    
    /// 更新元素
    func setByIndex(_ index: Int, element: E) {
        guard index >= 0 && index < self.size else {
            return
        }
        
        // 更新data里的元素
        self.data[index] = element
        // 更新tree里的元素
        self.setByTreeIndex(0, left: 0, right: self.size - 1, index: index, element: element)
    }
    
    private func setByTreeIndex(_ treeIndex: Int, left: Int, right: Int, index: Int, element: E) {
        if left == right {
            self.tree[treeIndex] = element
            return
        }
        
        let leftChildIndex = self.leftChildIndexBy(index: treeIndex)
        let rightChildIndex = self.rightChildIndexBy(index: treeIndex)
        
        let mid = (right - left) / 2 + left
        
        if index <= mid {
            self.setByTreeIndex(leftChildIndex, left: left, right: mid, index: index, element: element)
        }else if index > mid {
            self.setByTreeIndex(rightChildIndex, left: mid + 1, right: right, index: index, element: element)
        }
        self.tree[treeIndex] = self.merge(self.tree[leftChildIndex], self.tree[rightChildIndex])
    }
    
    
    private func leftChildIndexBy(index: Int) -> Int {
        return index * 2 + 1
    }
    
    private func rightChildIndexBy(index: Int) -> Int {
        return index * 2 + 2
    }
    
    
}
